/* -*- mode: c++; coding: utf-8-unix -*-
 */

#include "Algorithm.h"
#include "Log.h"
#include "mkstr"
#include <assert.h>
#include <math.h>
#include <string>

namespace alg
{

	using std::string;

	static const double halfpi = acos(0);
	static const double pi = halfpi*2;
	static const double twopi = pi*2;

	static const double EPSILON = 1e-15;

Workunit* Workunit::create(bool dft, bool center, const ChDef &src)
{
	if (dft) return new DFT_Workunit(center, src);
	else return new DCT_Workunit(center, src);
}

static inline img::ch_pixel_type getp(const ChDef &cd,
				      img::size_type m,
				      img::size_type n)
{
	assert(cd.image());
	return cd.image()->pixel(cd.channel(), m, n);
}
				      
       
static inline void setp(const ChDef &cd,
			img::size_type m,
			img::size_type n,
			img::ch_pixel_type v)
{
	if (cd.image())
		cd.image()->setPixel(cd.channel(), m, n, v);
}

void dct(const DCT_Workunit &wu)
{
	USELOG();

	typedef img::size_type ST;
	typedef double VT;

	const ST lastline = wu.startLine() + wu.numLines() - 1;
	LOG_DEBUG("DCT: Processing image lines " << wu.startLine() << ".." << lastline);

	const ChDef &src = wu.source();
	const ChDef &dst = wu.dest();

	const ST sM = src.image()->width(), sN = src.image()->height();
	const ST dM = dst.image()->width(), dN = dst.image()->height();
	const VT sq2 = sqrt(2);
	const VT C = 2/sqrt(sM*sN);
	const VT pi_p_2M = pi/(sM*2), pi_p_2N=pi/(sN*2);
	const ST shm = wu.center() ? -dM/2 : 0;
	const ST shn = wu.center() ? -dN/2 : 0;

	for (ST n=wu.startLine(); n<=lastline; n++)
	{
		const VT Fn = C*(n==0 ? 1/sq2 : 1);
		const VT pi_p_2N_n = pi_p_2N * (n+shn);

		for (ST m = 0; m < dM; m++)
		{
			const VT pi_p_2M_m = pi_p_2M * (m+shm);

			VT s = 0;
			for (ST u=0; u < sM; u++)
			{
				VT s2 = 0;
				for (ST v=0; v < sN; v++)
				{
					s2 += getp(src, u, v)
						* (cos(pi_p_2N_n*(2*v + 1)));
				}

				s+= s2*cos(pi_p_2M_m*(2*u + 1));
			}

			setp(dst, m, n, (Fn*(m==0 ? 1/sq2 : 1))*s);
		}
	}

	LOG_DEBUG("DCT finsihed processing lines " << wu.startLine() << ".." << lastline);
}

inline void A_SAMESIZE(const ChDef &a, const ChDef &b)
{
	if (a.image() && b.image())
	{
		if (!(a.image()->width() == b.image()->width() &&
		      a.image()->height() == b.image()->height()))
		{
			throw string("DFT requires that input "
				     "and output have the same size.");
		}
	}
}

inline bool zero(double x)
{
	return x < EPSILON;
}
inline double phase(double re, double im)
{
	if (zero(re))
	{
		if (zero(im)) return 0;
		if (im>0) return halfpi;
		if (im<0) return -halfpi;
	}
	const double atanimre = atan(im/re);
	if (re>0) return atanimre;
	if (zero(im) || im>=0) return atanimre + pi;
	if (im<0) return atanimre - pi;
}

void dft(const DFT_Workunit &wu)
{
	USELOG();

	typedef img::size_type ST;
	typedef double VT;

	const ST lastline = wu.startLine() + wu.numLines() - 1;
	LOG_DEBUG("DFT: Processing image lines " << wu.startLine() << ".." << lastline);

	const img::Channel sc = wu.source().channel();
	const ChDef
		&src = wu.source(),
		&Rdst = wu.realDest(),
		&Idst = wu.imagDest(),
		&Mdst = wu.magnitudeDest(),
		&Pdst = wu.phaseDest();

	assert(src.image());
	if (src.image()->width() != src.image()->height())
		throw string("DFT requires an image of size NxN");
	A_SAMESIZE(src, Rdst);
	A_SAMESIZE(src, Idst);
	A_SAMESIZE(src, Mdst);
	A_SAMESIZE(src, Pdst);

	const ST M = src.image()->width();
	const VT twopi_p_M = twopi/M;
	const VT invSqM = 1/sqrt(M);
	const VT invM = 1/M;
	const ST shm = wu.center() ? -M/2 : 0;

	for (ST n=wu.startLine(); n<=lastline; n++)
	{
		const ST nn = n+shm;
		
		for (ST m = 0; m<M; m++)
		{
			const ST mm = m+shm;
			
			VT res = 0, ims = 0;
			for (ST u = 0; u<M; u++)
			{
				const VT nn_u = nn*u;
				for (ST v = 0; v<M; v++)
				{
					const VT pv = getp(src, u, v);
					const VT arg = twopi_p_M * (nn_u + mm * v);
					res += pv*cos(arg);
					ims -= pv*sin(arg);
				}
			}

			setp(wu.realDest(), n, m, res);
			setp(wu.imagDest(), n, m, ims);
			setp(wu.magnitudeDest(), n, m,
			     sqrt(res*res + ims*ims));
			setp(wu.phaseDest(), n, m, phase(res, ims));
		}
	}

	LOG_DEBUG("DFT finsihed processing lines " << wu.startLine() << ".." << lastline);
}

std::string ChDef::toString() const
{
	return MKStr() << "ChDef [ Image=" << image()
		       << "; Channel=" << channel() << "]";
}
std::string Workunit::toString() const
{
	return MKStr() << "Workunit" << std::endl
		       << '[' << std::endl
		       << "  Src=<" << source().toString() << '>' << std::endl
		       << "  Range=(" << startLine() << "-->"
		       << startLine() + numLines() - 1 << ")" << std::endl
		       << ']';
}
std::string DFT_Workunit::toString() const
{
	return MKStr() << "DFT_Workunit" << std::endl
		       << '[' << std::endl
		       << "  Workunit     = <" << Workunit::toString() << '>' << std::endl
		       << "  RealDst = <" << realDest().toString() << '>' << std::endl
		       << "  ImagDst = <" << imagDest().toString() << '>' << std::endl
		       << ']';
}
std::string DCT_Workunit::toString() const
{
	return MKStr() << "DCT_Workunit" << std::endl
		       << '[' << std::endl
		       << "  Workunit = <" << Workunit::toString() << '>' << std::endl
		       << "  Dst = <" << dest().toString() << '>' << std::endl
		       << ']';
}

}
